import os
import shutil
import subprocess
import sys

import numpy as np
from pydub import AudioSegment

output_dir = "./wavs"
media_dir = "./medias"


# 采样率16k 保证和paddlespeech一致
def extract_audio(video_file_path):
    ext = os.path.splitext(video_file_path)[1].lower()
    output_path = video_file_path.replace(ext, ".wav")
    if os.path.exists(output_path):
        os.remove(output_path)
    # ffmpeg - i subtitle_demo1.mp4 - ac 1 - ar 16000 - vn input.wav
    # 定义 FFmpeg 命令
    ffmpeg_cmd = ['ffmpeg', '-i', video_file_path, '-ac', '1', '-ar', '16000', '-vn', output_path]
    # 执行 FFmpeg 命令
    subprocess.call(ffmpeg_cmd)
    return output_path


def convert_mp3_to_wav(mp3_file, wav_file):
    # 构建FFmpeg命令
    ffmpeg_cmd = [
        'ffmpeg',
        '-i', mp3_file,
        '-vn',  # 禁用视频流
        '-acodec', 'pcm_s16le',  # 设置音频编解码器为16位线性小端
        wav_file
    ]

    # 调用FFmpeg命令进行转换
    return subprocess.run(ffmpeg_cmd)


def is_audio_file(file_path):
    audio_extensions = ['.mp3', '.wav', '.ogg', '.opus', '.wma']
    ext = os.path.splitext(file_path)[1].lower()
    return ext in audio_extensions


def operate(audio_file):
    if not is_audio_file(audio_file):
        # 拿到新生成的音频的路径
        audio_file = extract_audio(audio_file)

    ext = os.path.splitext(audio_file)[1].lower()
    if ext != ".mp3" and ext != ".wav":
        return "仅支持mp3或wav格式！"
    if ext == ".mp3":
        # 拿到新生成的音频的路径
        dest = audio_file.replace('.mp3', '.wav')
        if os.path.exists(dest):
            os.remove(dest)
        try:
            sound = AudioSegment.from_mp3(audio_file)
            sound.set_channels(1)
            sound = sound.set_frame_rate(16000)
            sound = sound.set_channels(1)
            sound.export(dest, format="wav")
            audio_file = dest
        except Exception as e:
            print("转换失败：{}".format(e))
            return "转换失败：{}".format(e)

    split_audio(audio_file)
    return audio_file + "轉換完成"


import subprocess


def denoise_audio(input_file, output_file, denoise_value):
    # 构建FFmpeg命令
    ffmpeg_cmd = [
        'ffmpeg',
        '-i', input_file,
        '-vn',  # 禁用视频流
        '-af', 'anlmdn=d=0:i=-25:b=20',  # 使用非局部均值算法进行降噪，设置阈值为-25，步长为20
        '-y',  # 覆盖输出文件
        output_file
    ]

    # 调用FFmpeg命令进行降噪处理
    subprocess.run(ffmpeg_cmd)


def split_audio(audio_file_path):
    # 加载音频文件
    audio = AudioSegment.from_file(audio_file_path, format="mp3")

    # 获取音频总时长
    audio_time = len(audio)

    # 设置音频片段长度为60秒
    segment_length = 60 * 1000  # 毫秒

    file_name = os.path.split(audio_file_path)[1].replace(".wav", "")
    # 逐个切割音频文件为每60秒的小段
    for start_time in np.arange(0, audio_time, segment_length):
        stop_time = start_time + segment_length

        # 确保最后一段音频不超过总时长，如果有超过的部分，则直接取到文件末尾
        if stop_time > audio_time:
            stop_time = audio_time

        # 使用时间戳范围切割音频文件
        audio_chunk = audio[start_time:stop_time]

        # 保存新的音频文件
        new_file_name = output_dir + '/'+file_name + str(int(start_time / segment_length)) + '.wav'
        audio_chunk.export(new_file_name, format="wav")
        print('Finish splitting {} at [{},{}] ms'.format(audio_file_path, start_time, stop_time))


if __name__ == "__main__":
    if sys.platform.startswith('win'):
        media_dir = r"G:\workspace\python\snowboy-seasalt\seasalt\resources\sample1"
        output_dir = r"G:\workspace\python\snowboy-seasalt\seasalt\resources\sample2"
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    for media_file in os.listdir(media_dir):
        name,ext=os.path.splitext(media_file)
        input_path = os.path.join(media_dir, media_file)
        output_path = os.path.join(output_dir, media_file)
        if ext.lower().endswith(".wav"):
            shutil.copyfile(input_path,output_path)
        else:
            output_path= os.path.join(output_dir, name+".wav")
            # 读取MP3文件
            audio = AudioSegment.from_mp3(input_path)
            # 将MP3文件保存为WAV文件
            audio.export(output_path, format="wav")

